Add opt-in idle model unloading to free VRAM - #889
Conversation
Models load on demand at generate time but are only unloaded explicitly (POST /models/unload) or at shutdown, so an idle instance pins its weights indefinitely -- ~5.6GB for qwen-tts-1.7B, most of an 8GB card. Adds an asyncio reaper that unloads after a period of inactivity, gated on VOICEBOX_IDLE_UNLOAD_SECONDS and off by default. Idleness is read from the serial generation queue's own state so a model is never pulled out from under queued or in-flight work; the next request reloads it transparently. Refs jamiepine#595
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe backend adds an optional idle reaper that tracks generation activity and unloads TTS, Whisper, and LLM models after a configured idle period. It starts during application startup and resets its timer after each generation cycle. ChangesIdle model unloading
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant AppStartup
participant GenerationQueue
participant IdleReaper
participant ModelBackends
AppStartup->>GenerationQueue: init_queue()
AppStartup->>IdleReaper: start()
GenerationQueue->>IdleReaper: mark_activity() after each work cycle
IdleReaper->>GenerationQueue: check queued and running generations
IdleReaper->>ModelBackends: unload TTS, Whisper, and LLM after timeout
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/services/idle_unload.py`:
- Around line 76-82: Update the backend-loaded guard in _reaper so it considers
all backends cleared by _unload_all, including TTS, Whisper/STT, and LLM, before
skipping an unload cycle. Either remove the guard or reuse the corresponding
backend loaded-state checks, while preserving the existing exception-safe reaper
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6f4b4425-a983-4b08-8fd9-998861fa4425
📒 Files selected for processing (3)
backend/app.pybackend/services/idle_unload.pybackend/services/task_queue.py
Upstream voicebox never unloads models once loaded, so an idle instance pins ~4.3GB of GPU 1 indefinitely (most of the 2060 Super's 8GB). Adds VoiceboxConfig.build_ref (owner/repo@ref) so the service can build from a fork carrying the idle-unload patch, and defaults the timeout to 600s. Pointed at michael-borck/voicebox@feat/idle-model-unload pending jamiepine/voicebox#889; drop build_ref from puente.yml once that merges. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The reaper's guard only consulted the TTS backend, but _unload_all() frees Whisper and the LLM too. A session that dictates or chats without ever generating speech leaves TTS unloaded, so the guard short-circuited every cycle and those models stayed resident indefinitely — the exact leak this feature is meant to fix. Check all three backends, isolating each getter so one failing backend can't mask a loaded sibling. On error assume loaded rather than skipping: _unload_all() already isolates its own per-model failures, so attempting the unload is safer than pinning weights forever, and the exception is now logged instead of swallowed.
Closes #595 (the VRAM half of it — this adds the automatic trigger, not the manual load/unload UI button).
Problem
On a headless/server deployment there's no way for models to unload automatically. They load on demand at generate time, but the only things that ever unload them are an explicit
POST /models/unloador process shutdown (app.py). There's no idle path.So an instance that served one request an hour ago is still pinning its weights — ~5.6GB for qwen-tts-1.7B, which is most of an 8GB card. Anything else sharing that GPU is starved until you restart the container or remember to call the endpoint by hand.
Change
An asyncio reaper that unloads models after a period of inactivity.
VOICEBOX_IDLE_UNLOAD_SECONDS;0/unset disables it entirely, so existing behaviour is unchanged for anyone who doesn't opt in.tts/transcribe/llmunload calls already used on shutdown, with per-model failure isolation so one failure can't block the others._queued_generation_ids/_running_generation_tasks), not inferred externally, so a model is never pulled out from under a queued or running generation. A long generation continually re-arms the countdown.Three files: a new
backend/services/idle_unload.py, two lines inapp.pyto start it, and amark_activity()call in the generation worker'sfinallyblock (the one place every generation exits through, success/failure/cancel alike).Testing
Manual, on an RTX 2060 Super (8GB), Docker/CUDA, qwen-tts-1.7B:
Reclaims ~4.3GB of resident model weights. Verified separately that:
status=completed, no errors);I did not run Black across the touched files —
app.pyhas pre-existing formatting drift (theall_originsCORS block) and reformatting it would add unrelated churn to this diff. The new file and mytask_queue.pyedit are both Black-clean on their own.Disclosure: this patch was written with AI assistance (Claude), as was this description. It's tested end-to-end on my own build and the numbers above are real measurements from my hardware, not generated. If AI-assisted contributions aren't welcome here I understand — happy for this to be closed, and it stays a local patch on my build.
Summary by CodeRabbit